
#!/usr/bin/env python3
# -*- coding: utf-8 -*

# import matplotlib.pylab as plt
# import scipy.sparse as sparse
# import seaborn as sns
# import random
# import numpy as np
# import argparse
# import time
# import os
# from multiprocessing import Process
# from matirx_utility import read_data

# #TODO: Need to change.

# def my_pic_para():
#     my_pic = {}
#     my_pic["style"]         = "png"
#     my_pic["dpi"]           = 300
#     my_pic["myfontsize"]    = 12
#     my_pic["titlefontsize"] = 18
#     my_pic["figsize"] = [4,3]
#     my_pic["markersize"] = 0.5e-5
#     my_pic["transparent"] = False

#     # Solver01
#     # my_pic["save_pic_name"] = "solver01"
#     # my_pic["save_pic_dir"]  = "/home/matrix/solverchanllenge2021/solverchallenge21_01"
#     # my_pic["matrix_path"]   = "/home/matrix/solverchanllenge2021/solverchallenge21_01/jxpamg_matrix01"
#     # my_pic["save_matrix_path"]   = "/home/matrix/solverchanllenge2021/solverchallenge21_01/jxpamg_matrix01_change"
    
#     my_pic["save_pic_name"] = "mat_20X20"
#     my_pic["save_pic_dir"]  = "/home/dh/RAPdata"
#     my_pic["matrix_path"]   = "/home/dh/code/jxpamg/example/input/mat_20X20"
#     my_pic["save_matrix_path"]   = "/home/dh/RAPdata"

#     # my_pic["save_pic_name"] = "lixue"
#     # my_pic["save_pic_dir"]  = "/home/matrix/hsl_matrix"
#     # my_pic["matrix_path"]   = "/home/matrix/hsl_matrix/jxpamg_matrix"
#     # my_pic["save_matrix_path"]   = "/home/matrix/hsl_matrix/lixue_change"

#     # for test
#     # my_pic["save_pic_name"] = "test"
#     # my_pic["save_pic_dir"]  = "/home/mrz/code/python_script/pic_plot/plot_matrix"
#     # my_pic["matrix_path"]   = "/home/mrz/code/python_script/pic_plot/plot_matrix/mat_20X20"
#     # my_pic["save_matrix_path"]   = "/home/mrz/code/python_script/pic_plot/plot_matrix/mat_20X20_test"
    
#     my_pic["font"] = {'family': 'Times New Roman', 'size': '{fs}'.format(fs=my_pic["titlefontsize"])}
#     return my_pic

# def read_data(my_pic):
#     rptr = []
#     col = []
#     value = []
#     i     = 1
#     n     = 0
#     nnz   = 0
#     row_first = 0
#     col_first = 0
#     print("read matrix: \n" +my_pic["matrix_path"])
#     for line in open(my_pic["matrix_path"], "r" ):
#         line = line.strip("\n")
#         if(i == 1):
#             n = int(line)
#         elif(i > 1 and i <= n + 2):
#             if(i == 2):
#                 row_first = int(line)
#             elif(i == n + 2):
#                 nnz = int(line) - row_first
#             rptr.append(int(line) - row_first)
#         elif(i > n + 2 and i < n + 3 + nnz):
#             if(i == n + 3):
#                 col_first =int(line)
#             col.append(int(line) - col_first)
#         else:
#             value.append(float(line))
#             # print("%.15e" %(float(line)))
#         i = i + 1
#     print("read matrix success: n is {}, nnz is {}".format(n,nnz))

#     my_pic["csr"] = {}
#     my_pic["csr"]["n"] = n
#     my_pic["csr"]["nnz"] = nnz
#     my_pic["csr"]["rptr"] = rptr
#     my_pic["csr"]["col"]  = col
#     my_pic["csr"]["value"]= value

#     return my_pic

# def rearrange_matrix(my_pic):
#     # Note:
#     # 每行非零元数和元素值没有变动
#     # 变化的只有列标
#     n    = my_pic["csr"]["n"] 
#     # nnz  = my_pic["csr"]["n"] 
#     rptr = my_pic["csr"]["rptr"]
#     # col  = my_pic["csr"]["col"] 
#     my_order_tag = 0
#     rearrange_col = []
#     k = 0
#     #numpy
    
#     prepare_list = np.arange(n)
#     for i in range(n):
#         # T1 = time.perf_counter()
#         for j in range(rptr[i], rptr[i+1]):
#             if(i == j):
#                 my_order_tag = 1
#                 val_index =  j - rptr[i]
#         if(my_order_tag):
#             # T1 = time.perf_counter()
#             prepare_list = np.delete(prepare_list, i)
#             my_random = np.random.choice(prepare_list,rptr[i+1]-rptr[i] - 1)
#             my_random = np.insert(my_random,val_index,i)
#             prepare_list = np.insert(prepare_list,i,i)
#             # T2 = time.perf_counter()
#             # print('random和insert时间:%s ms' % ((T2 - T1)*1000))
#         else:
#             # T1 = time.perf_counter()
#             my_random = np.random.choice(prepare_list,rptr[i+1]-rptr[i])
#             # T2 = time.perf_counter()
#             # print('random时间:%s ms' % ((T2 - T1)*1000))
#         my_random = my_random.tolist()
#         rearrange_col = rearrange_col + my_random
#         if i == k*100000:
#             print("random row " + str(i) + " success! ", end = "\r")
#             k = k + 1
#         # print(my_random)
#         # print(col[rptr[i]:rptr[i+1]])
#         # T2 = time.perf_counter()
#         # print('a row time:%s ms' % ((T2 - T1)*1000))
#     # for i in range(n):
#     #     prepare_list = list(range(0,n))
#     #     del prepare_list[i]
#     #     for j in range(rptr[i], rptr[i+1]):
#     #         if(i == j):
#     #             my_order_tag = 1
#     #     if(my_order_tag):
#     #         my_random = random.sample(prepare_list,rptr[i+1]-rptr[i] - 1)
#     #         my_random.insert(0,i)
#     #     else:
#     #         my_random = random.sample(prepare_list,rptr[i+1]-rptr[i])
#     #     rearrange_col = rearrange_col + my_random
#     #     if i == k*100000:
#     #         print("random row " + str(i) + " success! ", end = "\r")
#     #         k = k + 1
#     print("success random! ")
#     # print(rearrange_col)
#     my_pic["csr"]["rearrange_col"] =  rearrange_col
#     print("rerrange matrix success")
#     return my_pic

# def save_rearrange_matrix(my_pic):
#     n    = my_pic["csr"]["n"]
#     nnz  = my_pic["csr"]["nnz"]
#     rptr = my_pic["csr"]["rptr"]
#     col  = my_pic["csr"]["col"]
#     value  = my_pic["csr"]["value"]

#     with open(my_pic["save_matrix_path"], "w+" ) as fo:
#         fo.write(str(n)+"\n")
#         for i in range(0, n + 1):
#             # fo.write(str(rptr[i])+"\n")
#             fo.write("%d" %(rptr[i])+"\n")
#         for i in range(0, nnz):
#             # fo.write(str(col[i])+"\n")
#             fo.write("%d" %(col[i])+"\n")
#         for i in range(0, nnz):
#             # fo.write(str(value[i])+"\n")
#             fo.write("%.15e" %(value[i])+"\n")

# def plot_my_pic(my_pic):
#     # arr = np.array([my_pic["csr"]["rptr"], my_pic["csr"]["rearrange_col"], my_pic["csr"]["value"]],
#     # dtype=object)
#     mat = sparse.csr_matrix((my_pic["csr"]["value"], my_pic["csr"]["col"], my_pic["csr"]["rptr"],), shape=(my_pic["csr"]["n"], my_pic["csr"]["n"]))
#     plt.rcParams['font.sans-serif']=[my_pic["font"]["family"]]
#     plt.figure(figsize=my_pic["figsize"],dpi=my_pic["dpi"])
#     plt.spy(mat, marker=',', markersize = my_pic["markersize"])
#     plt.savefig(my_pic["save_pic_dir"]+"/"+my_pic["save_pic_name"]+"_raw."+my_pic["style"],bbox_inches='tight')
#     print("save pic in " + my_pic["save_pic_dir"]+"/"+my_pic["save_pic_name"]+"_raw."+my_pic["style"])
#     plt.close()

# def plot_rearrange_pic(my_pic):
#     # arr = np.array([my_pic["csr"]["rptr"], my_pic["csr"]["rearrange_col"], my_pic["csr"]["value"]],
#     # dtype=object)
#     mat = sparse.csr_matrix((my_pic["csr"]["value"], my_pic["csr"]["rearrange_col"], my_pic["csr"]["rptr"],), shape=(my_pic["csr"]["n"], my_pic["csr"]["n"]))
#     plt.rcParams['font.sans-serif']=[my_pic["font"]["family"]]
#     plt.figure(figsize=my_pic["figsize"], dpi=my_pic["dpi"])
#     plt.spy(mat, markersize = my_pic["markersize"])
#     plt.savefig(my_pic["save_pic_dir"]+"/"+my_pic["save_pic_name"]+"."+my_pic["style"],bbox_inches='tight')
#     print("save pic in " + my_pic["save_pic_dir"]+"/"+my_pic["save_pic_name"]+"."+my_pic["style"])
#     plt.close()

# def matrix_bandwith(mat):
#     bandwidth = 0
#     for i in np.arange(mat.nnz):
#         bandwidth = max(bandwidth, abs(mat.row[i] - mat.col[i]))
#     return bandwidth

# # def matrix_bandwith_definition(mat):
# #     bandwidth = 0
# #     for i in np.arange(mat.nnz):
# #         bandwidth = max(bandwidth, abs(mat.row[i] - mat.col[i]))
# #     return bandwidth

# def get_mat_feature(mat):
#     #matrix_bandwidth = bandwidth(mat.todense())
#     bandwidth = matrix_bandwith(mat)
#     sparsity = float(mat.nnz) / float(mat.shape[0]*mat.shape[1])
#     nnz_per_row = mat.nnz / mat.shape[0]
#     print("=====================")
#     print(">> matrix feature is:")
#     print(">> bandwith upper is {0}".format(str(bandwidth)))
#     print(">> sparsity is {0:.3e}".format(sparsity))
#     print(">> nnz per row is {0}".format(str(nnz_per_row)))
#     print("=====================")



# def trans2coo(my_pic):
#     mat = sparse.csr_matrix((my_pic["csr"]["value"], my_pic["csr"]["col"], my_pic["csr"]["rptr"],), shape=(my_pic["csr"]["n"], my_pic["csr"]["n"]))
#     # my_pic["csr_mat"] = mat
#     my_pic["coo_mat"] = mat.tocoo()
#     return my_pic

# #TODO:画图代码不够美观
# #COPY石油大学的部分代码
# def plot_heatmap(my_pic):
#     # arr = np.array([my_pic["csr"]["rptr"], my_pic["csr"]["rearrange_col"], my_pic["csr"]["value"]],
#     # dtype=object)
#     mat = my_pic["coo_mat"]
#     plt.rcParams['font.sans-serif']=[my_pic["font"]["family"]]
#     fig = plt.figure(figsize=my_pic["figsize"], dpi = my_pic["dpi"])
#     ax = plt.gca()      #获取到当前坐标轴信息
#     ax.xaxis.set_ticks_position('top')   #将X坐标轴移到上面
#     ax.invert_yaxis()   #反转Y坐标轴
#     plt.scatter(mat.row, mat.col,c = mat.data, s = my_pic["markersize"])
#     plt.colorbar()
#     plt.xlim(0,mat.shape[1])
#     plt.ylim(mat.shape[0],0) 
#     fig.savefig("{0}/{1}.{2}".format(my_pic["save_pic_dir"],my_pic["save_pic_name"],my_pic["style"]),transparent=my_pic["transparent"],bbox_inches='tight')
#     plt.close(fig)
#     print("save pic in " + my_pic["save_pic_dir"]+"/"+my_pic["save_pic_name"]+"."+my_pic["style"])

# def main():
#     my_pic = my_pic_para()
#     my_pic = read_data(my_pic)
#     my_pic = trans2coo(my_pic)
#     plot_heatmap(my_pic)
#     # get_mat_feature(my_pic["coo_mat"])
#     #plot_my_pic(my_pic)

#     # rerange matrix
#     # my_pic = rearrange_matrix(my_pic)
#     # save_rearrange_matrix(my_pic)
#     # plot_rearrange_pic(my_pic)



import scipy.sparse as sparse
from scipy.sparse import csr_matrix, coo_matrix, load_npz, save_npz
from scipy.sparse.linalg import eigs
import time
import sys
import pandas as pd
import numpy as np
import os
import json
import argparse
#+++++++only for test++++++++++
import matplotlib.pyplot as plt
#++++++++++++++++++++++++++++++
import seaborn as sns

from multiprocessing import Pool
from multiprocessing import Manager

from matirx_utility import read_data

def make_parser(): 
    parser = argparse.ArgumentParser(
        prog="Compute matrix feature",
        description="Compute matrix feature.")
    parser.add_argument(
            "MatDIR",
            metavar="MatrixDIR",
            help="Destination matrix files."
        )
    parser.add_argument(
            "SaveDIR",
            metavar="SaveFeatureDIR",
            help="Destination Save matrix feature."
        )
    return parser

def plot_matrix(matrix, title = "", file = "matrix", save_path = "./"):
    """
    Plots a matrix.
    """
    save_path = os.path.abspath(save_path)
    plt.figure(figsize=(20, 20))
    plt.title(title + "_pattern")
    plt.spy(matrix, markersize=0.02)
    # plt.show()
    file_path = save_path + "/" + file + ".png"
    #file_path = "plot/" + file + "_%d.png" % refine
    plt.savefig(file_path, dpi = 500, bbox_inches='tight')

def plot_matrix_data(val, refine, title="", file="matrix"):
    """
    Plots matrix data.
    """
    plt.figure(figsize=(20, 20))
    sns.histplot(data = val, kde=False, discrete=True, binwidth=1, stat="percent")
    plt.xlabel("nnz per row")
    plt.title(title + " (refine = %d)" % refine)
    #file_path = "plot/" + file + "_%d.png" % refine
    file_path = "./" + file + "_%d.png" % refine
    plt.savefig(file_path, dpi = 500, bbox_inches='tight')

def test():
    parser = make_parser()
    args = parser.parse_args()
    #print(args.MatDIR)
    my_matrix = read_data(args.MatDIR)
    filename = os.path.split(args.MatDIR)[1]
    plot_matrix(my_matrix, filename, filename, args.SaveDIR)
    #plot_matrix_data()

if __name__=="__main__":
    test()
    #main()
